import torch
import torch.nn as nn

def triplet_margin_loss_hard_neg(
    input_a, input_b, mtc_fn
):
    b_a, b_b = input_a.size(0), input_b.size(0)

    matrix = [[] for i in range(b_a)]

    for i in range(b_a):
        m = matrix[i]
        a = input_a[i].unsqueeze(0)

        for j in range(b_b):
            b = input_b[j].unsqueeze(0)
            m.append(mtc_fn(a, b))


    nums = min(b_a, b_b)
    samples = [[] for i in range(nums)]

    for n in range(nums):
        position = matrix[n][n]

        neg = []
        for i in range(b_a):
            if i == n:
                continue

                neg.append(matrix[i][n])

        for j in range(b_b):
                if n == j:
                    continue

                neg.append(matrix[n][j])
        
        if len(neg) == 0:
            negtive = 0
        else:
            negtive = min(neg)
        samples[n].extend([position, negtive])

    loss = sum(list(map(
        lambda x : max(torch.tensor(0., requires_grad=True), 1+x[0]-x[1]), 
        samples
    )))
    loss /= nums

    return loss


